from zhipuai import ZhipuAI
#该类的作用就是通过一个prompt，得到答案

def get_ans(prompt):
    # 需要修改为自己的智谱的api_key
    #使用智谱线上的大模型
    client = ZhipuAI(api_key="your_api_key")
    response = client.chat.completions.create(
        model="glm-4",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],
        top_p=0.3,
        temperature=0.45,
        max_tokens=1024,
        stream=True,
    )
    #得到答案 response是一个对象

    ans = ""
    for trunk in response:
        #trunk类似于一个字一个字预测
        ans += trunk.choices[0].delta.content
    #ans为回复
    ans = ans.replace("\n\n", "\n")
    return ans


if __name__ == "__main__":
    prompt = "什么是大模型"
    ans = get_ans(prompt)
    print(ans)
